package util

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"github.com/pkg/sftp"
	"golang.org/x/crypto/ssh"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/user"
	"path"
	"time"
)

type Cli struct {
	IP         string      //IP地址
	Username   string      //用户名
	Password   string      //密码
	PrivateKey string      //秘钥登录
	Port       int         //端口号
	client     *ssh.Client //ssh客户端
	LastResult string      //最近一次Run的结果
	SftpClient *sftp.Client
}

//创建命令行对象
//@param ip IP地址
//@param username 用户名
//@param password 密码
//@param port 端口号,默认22
func NewSsh(ip string, username string, password string, publicKey string, port ...int) *Cli {
	cli := new(Cli)
	cli.IP = ip
	cli.Username = username
	cli.Password = password
	cli.PrivateKey = publicKey
	if len(port) <= 0 {
		cli.Port = 22
	} else {
		cli.Port = port[0]
	}
	return cli
}

//执行shell
//@param shell shell脚本命令
func (c Cli) Run(shell string) (string, error) {
	if c.client == nil {
		if err := c.connect(); err != nil {
			return "", err
		}
	}
	session, err := c.client.NewSession()
	if err != nil {
		return "", err
	}
	defer session.Close()
	buf, err := session.CombinedOutput(shell)

	c.LastResult = string(buf)
	return c.LastResult, err
}

func (c Cli) RunMonitor(shell string, queryLine QueryLine) error {
	if c.client == nil {
		if err := c.connect(); err != nil {
			return err
		}
	}
	session, err := c.client.NewSession()
	if err != nil {
		return err
	}
	defer session.Close()

	stdoutRs, err := session.StdoutPipe()
	var stderr bytes.Buffer
	session.Stderr = &stderr
	readerRs := bufio.NewReader(stdoutRs)
	if err := session.Start(shell); err != nil {
		return err
	}

	errChan := make(chan error)
	go func() {
		errChan <- session.Wait()
		defer close(errChan)
	}()

	go func() {
		for {
			line, isPrefix, errLine := readerRs.ReadLine()
			if !isPrefix {
				if errLine != nil && errLine == io.EOF {
					break
				} else if errLine != nil {
					return
				}
				queryLine(string(line))
			}
		}
	}()

	var err2 error
	Timeout := 3600 * time.Second
	select {
	case <-time.After(Timeout):
		err2 = errors.New(fmt.Sprintf("cmd run timeout, cmd [%s], time[%v]", shell, Timeout))
	case err2 = <-errChan:
	}
	return err2
}

//本地上传文件
func (c *Cli) UploadFile(localFilePath string, remotePath string, queryLine QueryLine) error {
	if err := c.connect(); err != nil {
		return errors.New("ssh 登录失败 :" + err.Error())
	}

	//打开本地文件流
	srcFile, err := os.Open(localFilePath)
	if err != nil {
		return errors.New(fmt.Sprintln("os.Open error : ", localFilePath, err.Error()))
	}
	//关闭文件流
	defer func() {
		srcFile.Close()
		c.SftpClient.Close()
	}()

	//上传到远端服务器的文件名,与本地路径末尾相同
	var remoteFileName = path.Base(localFilePath)
	//打开远程文件,如果不存在就创建一个
	dstFile, err := c.SftpClient.Create(path.Join(remotePath, remoteFileName))
	if err != nil {
		return errors.New(fmt.Sprintln("sftpClient.Create error : ", path.Join(remotePath, remoteFileName), err.Error()))
	}
	//关闭远程文件
	defer func() {
		dstFile.Close()
	}()

	fileStat, _ := srcFile.Stat()
	queryLine("文件大小：" + FormatFileSize(fileStat.Size()))

	fileInfo, err := ioutil.ReadAll(srcFile)
	if err != nil {
		return err
	}
	dstFile.Write(fileInfo)

	return nil
}

//远程写入文件
func (c *Cli) RemoteWriteString(data, remotePath string) error {
	if err := c.connect(); err != nil {
		return errors.New("ssh 登录失败 :" + err.Error())
	}
	//打开远程文件,如果不存在就创建一个
	_ = c.SftpClient.Remove(remotePath)
	dstFile, err := c.SftpClient.Create(remotePath)
	if err != nil {
		return errors.New(fmt.Sprintln("sftpClient.Create error : ", remotePath, err.Error()))
	}
	//关闭远程文件
	defer func() {
		c.SftpClient.Close()
		dstFile.Close()
	}()
	data = `#!/bin/bash
			set -e	
			` + data
	_, err = dstFile.Write([]byte(data))
	if err != nil {
		return errors.New(fmt.Sprintln("sftpClient.Write error : ", remotePath, err.Error()))
	}
	return nil
}

//连接
func (c *Cli) connect() error {
	config := ssh.ClientConfig{}
	if c.PrivateKey != "0" {
		currentUser, _ := user.Current()
		id_rsa_data, err := ioutil.ReadFile(currentUser.HomeDir + "/.ssh/id_rsa")
		if err != nil {
			return err
		}
		signer, err := ssh.ParsePrivateKey(id_rsa_data)
		if err != nil {
			return err
		}
		config = ssh.ClientConfig{
			User: c.Username,
			Auth: []ssh.AuthMethod{
				ssh.PublicKeys(signer),
			},
			HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
				return nil
			},
			Timeout: 10 * time.Second,
		}
	} else {
		config = ssh.ClientConfig{
			User: c.Username,
			Auth: []ssh.AuthMethod{ssh.Password(c.Password)},
			HostKeyCallback: func(hostname string, remote net.Addr, key ssh.PublicKey) error {
				return nil
			},
			Timeout: 10 * time.Second,
		}
	}

	addr := fmt.Sprintf("%s:%d", c.IP, c.Port)
	sshClient, err := DialTimeout("tcp", addr, &config)
	//sshClient, err := ssh.Dial("tcp", addr, &config)
	if err != nil {
		return err
	}
	c.client = sshClient

	// create sftp client
	sftpClient, err := sftp.NewClient(sshClient)
	if err != nil {
		return err
	}
	c.SftpClient = sftpClient
	return nil
}

func DialTimeout(network, addr string, config *ssh.ClientConfig) (*ssh.Client, error) {
	conn, err := net.DialTimeout(network, addr, config.Timeout)
	if err != nil {
		return nil, err
	}
	// set timeout for connection
	timeFactor := time.Duration(3)
	err = conn.SetDeadline(time.Now().Add(config.Timeout * timeFactor))
	if err != nil {
		conn.Close()
		return nil, err
	}

	// do SSH handshake
	c, chans, reqs, err := ssh.NewClientConn(conn, addr, config)
	if err != nil {
		return nil, err
	}

	// cancel connection read/write timeout
	err = conn.SetDeadline(time.Time{})
	if err != nil {
		conn.Close()
		return nil, err
	}
	return ssh.NewClient(c, chans, reqs), nil
}
